package brave.kafka.clients;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.internal.Nullable;
import brave.messaging.MessagingRequest;
import brave.propagation.TraceContext;
import brave.sampler.SamplerFunction;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;

import java.util.stream.Stream;

/**
 * 接收后跟踪处理器
 * 仅针对于 kafkaTracing.nextSpan(record) 默认实现方法进行增强 ，参数必须是 ConsumerRecord
 * 当参数 List<ConsumerRecord> 复数行为的时候，放弃绑定，业务自行解析处理
 *
 * @author zsy
 * @see Tracing
 * @see brave.kafka.clients.TracingConsumer
 * @since 2023/11/20
 */
@Slf4j
public class TracingKafkaPostProcessor {

    final ThreadLocal<SpanScope> local;
    final KafkaTracing kafkaTracing;
    final Tracing tracing;
    final Tracer tracer;
    final TraceContext.Extractor<KafkaConsumerRequest> extractor;
    final TraceContext.Injector<KafkaConsumerRequest> injector;
    final SamplerFunction<MessagingRequest> sampler;
    /**
     * 是否通过参数采集
     * 如果为 true，将遍历
     */
    final boolean parameterAcquisition;
    @Nullable
    final String remoteServiceName;

    public TracingKafkaPostProcessor(KafkaTracing kafkaTracing, boolean parameterAcquisition) {
        this.kafkaTracing = kafkaTracing;
        this.tracing = kafkaTracing.messagingTracing.tracing();
        this.tracer = tracing.tracer();
        this.extractor = kafkaTracing.consumerExtractor;
        this.sampler = kafkaTracing.consumerSampler;
        this.injector = kafkaTracing.consumerInjector;
        this.remoteServiceName = kafkaTracing.remoteServiceName;
        this.parameterAcquisition = parameterAcquisition;
        this.local = new ThreadLocal<>();
    }

    public void nextSpan(ProceedingJoinPoint pjp) {
        if (parameterAcquisition) {
            set(Stream.of(pjp.getArgs())
                    .filter(arg -> arg instanceof ConsumerRecord<?, ?>)
                    .map(ConsumerRecord.class::cast)
                    .findFirst()
                    .map(record -> setConsumerSpan(kafkaTracing.nextSpan(record), record.topic(), pjp.getSignature()))
                    .orElseGet(() -> setConsumerSpan(tracer.nextSpan(), null, pjp.getSignature())));
        } else {
            set(setConsumerSpan(tracer.nextSpan(), null, pjp.getSignature()));
        }
    }

    Span setConsumerSpan(Span span, String topic, Signature signature) {
        if (!span.isNoop()) {
            span.name("next-message").kind(Span.Kind.CONSUMER);
            span.tag("class", signature.getDeclaringTypeName());
            span.tag("method", signature.getName());
            if (topic != null) span.tag(KafkaTags.KAFKA_TOPIC_TAG, topic);
            if (remoteServiceName != null) span.remoteServiceName(remoteServiceName);
            span.start();
        }
        return span;
    }

    /**
     * 参考rabbit brave释放过程
     * Throwable error = null;
     * try {
     * return message;
     * } catch (Throwable t) {
     * error = t;
     * throw t;
     * } finally {
     * if (error != null) listenerSpan.error(error);
     * listenerSpan.finish();
     * ws.close();
     * }
     *
     * @param spanScope
     * @param error
     */
    protected void release(SpanScope spanScope, Throwable error) {
        if (spanScope != null) {
            if (error != null) spanScope.span.error(error);
            spanScope.span.finish();
            spanScope.scope.close();
        }
    }

    protected void set(Span span) {
        release(local.get(), null);
        local.set(new SpanScope(span, tracer.withSpanInScope(span)));
    }

    public void close(Throwable error) {
        release(local.get(), error);
        local.remove();
    }

    public record SpanScope(Span span, Tracer.SpanInScope scope) {
    }

}
